package ffmpeg

/*
#cgo LDFLAGS: -lavformat -lavcodec -lswresample -lavutil -lz -lm

#include <stdlib.h>

#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
*/
import "C"

import (
	"errors"
	"fmt"
	"log"
	"unsafe"
)

type Output struct {
	url         string
	fmtCtx      *C.AVFormatContext
	audioStream *AVStream
	videoStream *AVStream
}

func Create(url, format string) (*Output, error) {
	csUrl := C.CString(url)
	defer C.free(unsafe.Pointer(csUrl))
	var csFormat *C.char = nil
	if format != "" {
		csFormat = C.CString(format)
	}
	defer C.free(unsafe.Pointer(csFormat))
	var fmtCtx *C.AVFormatContext
	if retCode := C.avformat_alloc_output_context2(&fmtCtx, nil, csFormat, csUrl); retCode < 0 {
		return nil, errors.New(fmt.Sprintf("Call to avformat_alloc_output_context2(%s) failed: %d.", url, retCode))
	}
	return &Output{url, fmtCtx, nil, nil}, nil
}

func (f *Output) Write(in *Input, startInSec, durationInSec int) error {
	for _, istream := range in.Streams {
		if istream.codec.codec_type == C.AVMEDIA_TYPE_VIDEO {
			if f.videoStream == nil {
				f.videoStream = (*AVStream)(C.avformat_new_stream(f.fmtCtx, nil))
				if f.videoStream == nil {
					return errors.New("Failed to allocate stream!")
				}
				if retCode := C.avcodec_copy_context(f.videoStream.codec, istream.codec); retCode < 0 {
					return errors.New("Failed to copy codec context!")
				}
				if f.fmtCtx.oformat.flags&C.AVFMT_GLOBALHEADER != 0 {
					f.videoStream.codec.flags |= C.CODEC_FLAG_GLOBAL_HEADER
				}
			}
		} else if istream.codec.codec_type == C.AVMEDIA_TYPE_AUDIO {
			if f.audioStream == nil {
				f.audioStream = (*AVStream)(C.avformat_new_stream(f.fmtCtx, nil))
				if f.audioStream == nil {
					return errors.New("Failed to allocate stream!")
				}
				if retCode := C.avcodec_copy_context(f.audioStream.codec, istream.codec); retCode < 0 {
					return errors.New("Failed to copy codec context!")
				}
				if f.fmtCtx.oformat.flags&C.AVFMT_GLOBALHEADER != 0 {
					f.videoStream.codec.flags |= C.CODEC_FLAG_GLOBAL_HEADER
				}
			}
		} else {
			log.Printf("Ignore stream type: %s.\n", istream.MediaType())
		}
	}

	if f.fmtCtx.oformat.flags&C.AVFMT_NOFILE == 0 {
		if retCode := C.avio_open(&f.fmtCtx.pb, (*_Ctype_char)(unsafe.Pointer(&f.fmtCtx.filename)), C.AVIO_FLAG_WRITE); retCode < 0 {
			return errors.New(fmt.Sprintf("Failed to open '%s' for write: %d.", f.url, retCode))
		}
	}

	if retCode := C.avformat_write_header(f.fmtCtx, nil); retCode < 0 {
		return errors.New(fmt.Sprintf("Failed to write header for '%s': %d.", f.url, retCode))
	}

	for {
		packet := C.AVPacket{data: nil, size: 0}
		if retCode := C.av_read_frame(in.fmtCtx, &packet); retCode < 0 {
			break
		}
		defer C.av_free_packet(&packet)

		inStream := in.Streams[packet.stream_index]
		codecType := inStream.codec.codec_type

		ptsInSec := int64(packet.pts) * int64(inStream.time_base.num) / int64(inStream.time_base.den)
		if ptsInSec < int64(startInSec) {
			continue
		}

		endInSec := (int64(packet.pts) + int64(packet.duration)) * int64(inStream.time_base.num) / int64(inStream.time_base.den)
		shouldStop := durationInSec > 0 && endInSec >= int64(startInSec+durationInSec)

		var outStream *AVStream
		if codecType == C.AVMEDIA_TYPE_VIDEO {
			outStream = f.videoStream
		} else if codecType == C.AVMEDIA_TYPE_AUDIO {
			outStream = f.audioStream
		} else {
			continue
		}
		packet.stream_index = outStream.index

		C.av_packet_rescale_ts(&packet, inStream.time_base, outStream.time_base)
		if retCode := C.av_interleaved_write_frame(f.fmtCtx, &packet); retCode < 0 {
			return errors.New(fmt.Sprintf("Failed to write frame: %d.", retCode))
		}
		if shouldStop {
			break
		}
	}
	return nil
}

func (f *Output) Close() error {
	if f == nil || f.fmtCtx == nil {
		return errors.New("File never opened!")
	}
	streams := getStreams(f.fmtCtx)
	for _, stream := range streams {
		if retCode := C.avcodec_close(stream.codec); retCode < 0 {
			log.Printf("Call to avcodec_close for stream %d failed: %d\n", stream.index, retCode)
		}
	}
	if f.fmtCtx.oformat.flags&C.AVFMT_NOFILE == 0 {
		C.avio_close(f.fmtCtx.pb)
	}
	C.avformat_free_context(f.fmtCtx)
	return nil
}
